import os
import os.path as osp
import base64
import io
import time
from PIL import Image
import numpy as np
import cv2
import json
import requests

from rec.ocr  import ocr_rec_func
from rec.yolov8_api  import api_yolov8_det_img, pack_yolov8_handwrite_data, pack_yolov8_pagenum_data


from collections import defaultdict
from flask import Flask
from flask import request

app = Flask(__name__)

DEBUG_MODEL = 0
seal_url = "http://183.221.0.158:29986/api/algo/docLayout"
pagenum_url = "http://183.221.0.158:29986/api/algo/pageNum"


@app.route('/hi', methods=['GET', 'POST'])
def home():
    return '<h1>Home</h1>'

@app.route('/api/ocr/recognize', methods=['POST'])
def api_ocr_recognize():
    json_data = request.get_json()
    if type(json_data) != dict: 
        print("type error...", type(json_data), " != dict")
        return json.dumps({"code":1201, "msg":"参数错误"})

    img_b64 = json_data.get("img_base64", None)
    if img_b64 is None:
        return json.dumps({"code":1202, "msg":"参数名错误"})
        
    base64_data = base64.b64decode(img_b64)
    imageBuf = np.fromstring(base64_data, dtype=np.uint8)
    bgr_img = cv2.imdecode(imageBuf, cv2.IMREAD_COLOR)
    print(bgr_img.shape)
    if len(bgr_img.shape) != 3:
        return json.dumps({"code":1203, "msg":"图片非3通道错误"})
        
    categories = {}
    categories["title"] = ""
    categories["seals"] = []
    categories["handwritten_signatures"] = []
    categories["signature_dates"] = []
    categories["page_number"] = {}
    

    ret_data_handwrite = api_yolov8_det_img(bgr_img, img_b64, filter_score=1, url=seal_url, debug=DEBUG_MODEL)
    if ret_data_handwrite["code"] == 0:
        pack_yolov8_handwrite_data(categories, bgr_img, ret_data_handwrite, DEBUG_MODEL)

    ret_data_pagenum = api_yolov8_det_img(bgr_img, img_b64, filter_score=1, url=pagenum_url, debug=DEBUG_MODEL)
    if ret_data_pagenum["code"] == 0:
        pack_yolov8_pagenum_data(categories, bgr_img, ret_data_pagenum, DEBUG_MODEL)

    ocr_text = ocr_rec_func(bgr_img, DEBUG_MODEL)

    ret_data = {}
    ret_data["code"] = 0
    ret_data["msg"] = "ok"
    item_data = {}
    item_data["categories"] = categories
    item_data["ocr_text"] = ocr_text
    ret_data["data"] = item_data
    ret_json_str = json.dumps(ret_data, ensure_ascii=False, indent=4)
    return ret_json_str

@app.route('/api/ocr/text_rec', methods=['POST'])
def api_ocr_text_rec():
    json_data = request.get_json()
    if type(json_data) != dict: 
        print("type error...", type(json_data), " != dict")
        return json.dumps({"code":1201, "msg":"参数错误"})

    img_b64 = json_data.get("img_base64", None)
    if img_b64 is None:
        return json.dumps({"code":1202, "msg":"参数名错误"})
        
    base64_data = base64.b64decode(img_b64)
    imageBuf = np.fromstring(base64_data, dtype=np.uint8)
    bgr_img = cv2.imdecode(imageBuf, cv2.IMREAD_COLOR)
    print(bgr_img.shape)
    if len(bgr_img.shape) != 3:
        return json.dumps({"code":1203, "msg":"图片非3通道错误"})
        
    # categories = {}
    # categories["title"] = ""
    # categories["seals"] = []
    # categories["handwritten_signatures"] = []
    # categories["signature_dates"] = []
    # categories["page_number"] = {}
    

    ocr_text = ocr_rec_func(bgr_img, DEBUG_MODEL)

    ret_data = {}
    ret_data["code"] = 0
    ret_data["msg"] = "ok"
    item_data = {}
    # item_data["categories"] = categories
    item_data["ocr_text"] = ocr_text
    ret_data["data"] = item_data
    ret_json_str = json.dumps(ret_data, ensure_ascii=False, indent=4)
    return ret_json_str




if __name__ == '__main__':
    app.run("0.0.0.0", 9966)
